use std::collections::HashSet;
use std::collections::VecDeque;

type Position = (i32, i32);

fn parse_input() -> Vec<Vec<u32>> {
    let input = std::fs::read_to_string("./day09/input.txt").unwrap();
    input.lines().map(|line| -> Vec<u32> {
        line.chars().map(|c| c.to_digit(10).unwrap()).collect()
    }).collect()
}

fn adjacents((x, y): Position) -> [Position; 4] {
    [(x+1, y), (x-1, y), (x, y-1), (x, y+1)]
}

#[allow(unused)]
fn part1(heights: &Vec<Vec<u32>>) {
    let row = heights.len() as i32;
    let col = heights[0].len() as i32;

    let height = |(x, y): Position| -> Option<u32> {
        if x >= 0 && x < col && y >= 0 && y < row { 
            Some(heights[y as usize][x as usize])
        } else {
            None
        }
    };

    let mut result = 0;
    for y in 0..row {
        for x in 0..col {
            let curr_height = height((x, y)).unwrap();
            if adjacents((x, y)).into_iter()
                                     .map(|pos| height(pos))
                                     .filter(|h| h.is_some())
                                     .map(|h| h.unwrap())
                                     .all(|h| h > curr_height)
            {
                result += (1 + curr_height);
            }
        }
    }

    println!("Part1: {}", result);
}   

#[allow(unused)]
fn part2(heights: &Vec<Vec<u32>>) {
    let row = heights.len() as i32;
    let col = heights[0].len() as i32;
    let mut touched: Vec<Vec<bool>> = vec![vec![false; col as usize]; row as usize];

    let height = |&(x, y): &Position| -> Option<u32> {
        if x >= 0 && x < col && y >= 0 && y < row { 
            Some(heights[y as usize][x as usize])
        } else {
            None
        }
    };

    let mut expend_basin = |(x, y): Position| -> usize {
        if touched[y as usize][x as usize] == true {
            return 0;
        }
        let mut queue: VecDeque<Position> = VecDeque::new();
        let mut basins: HashSet<Position> = HashSet::new();
        queue.push_back((x, y));

        loop {
            match queue.pop_front() {
                None => break,
                Some(pos) => {
                    adjacents(pos).into_iter()
                        .filter(|adj_pos: &(i32, i32)| -> bool {
                            match height(&adj_pos) {
                                None => false,
                                Some(h) => h != 9,
                            }
                        })
                        .for_each(|adj_pos| { 
                            if !basins.contains(&adj_pos) {
                                let (x, y) = adj_pos;
                                basins.insert(adj_pos.clone());
                                queue.push_back(adj_pos);
                                touched[y as usize][x as usize] = true;
                            }
                        });
                }
            }
        }

        basins.len()
    };

    let mut basin_size = Vec::new();
    for y in 0..row {
        for x in 0..col {
            if height(&(x, y)).unwrap() == 9 {
                continue;
            }
            basin_size.push(expend_basin((x, y)));
        }
    }    

    basin_size.sort_by(|a, b| b.cmp(a));
    println!("{:?}", basin_size);
    let results: Vec<usize> = basin_size.into_iter().take(3).collect();

    println!("Part2: {}", results[0] * results[1] * results[2]);

}

fn main() {
    let heights = parse_input();
    // part1(&heights);
    part2(&heights);
}
